

from envs import next_step_back_mp_coo_large_scale_twice_timeout_noprofiling, next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling

import random
import numpy as np


def setup_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
     
setup_seed(20230430)


# sparse_name_list = ['gen2', '176bit', 'bcsstk13', 'biplane-9', 
#                     'ch7-8-b3', 'luxembourg_osm', 'cond-mat-2003', 'astro-ph', 
#                     '3D_28984_Tetra', 'rajat25', 'c-69', 'rajat16', 
#                     'net4-1', 'tsyl201', 'matrix-new_3', 'ct20stif']

sparse_name_list = ['t2d_q9', 'fv2', 'fv3', 'big_dual', 
                    'wing', 'TSOPF_RS_b39_c7', 'airfoil_2d', 'jan99jac120sc', 
                    'Trec13', 'Zd_Jac2_db', 'bibd_17_8', 'image_interp', 
                    'stat96v2', 'belgium_osm', 'Chevron3', 'parabolic_fem']

rank = 0
world_size = 2

block_num, thread_num, element_num = 512, 256, 2

# sparse_name_list = ['ch7-8-b3']
# sparse_name = 'ch7-8-b3'

mean_gflops_list = []
std_gflops_list = []

for sparse_name in sparse_name_list:
    try:
        init_gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, block_num, thread_num, element_num, rank, world_size)

    except Exception as e:
        
        flag_continue = True 
        while flag_continue:
            try:
                init_gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, block_num, thread_num, element_num, rank, world_size) # get next state

                flag_continue = False
            except:
                print('large: ', sparse_name, " ", block_num, " ", thread_num, " ", element_num)

    print('first: ', init_gflops)

    parameters_scale = np.array([(128, 4096, 128), (32, 1024, 32), (1, 32, 1)])

    gflops_list = []
    parameters_num = 32
    program_feature = [0, 0, 0]

    for i in range(30):
        
        for j in range(3):    
            policy_select_value = random.randint(0, parameters_num-1)  
            program_feature[j] = policy_select_value * parameters_scale[j][2] \
                                    + parameters_scale[j][0] # get
                                    
        try:
            gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                program_feature[2], rank, world_size)
        except Exception as e:
            
            flag_continue = True 
            while flag_continue:
                try:
                    gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(sparse_name, program_feature[0], program_feature[1], 
                                                                                                    program_feature[2], rank, world_size) # get next state
                    flag_continue = False
                except:
                    print('large: ', sparse_name, " ", program_feature[0], program_feature[1],  program_feature[2],)
        gflops_list.append(gflops)
        print('config: ', program_feature)
        
    gflops_list = np.array(gflops_list)
    
    # print(sparse_name)
    # print(np.mean(gflops_list))
    # print(max(np.abs(gflops_list - init_gflops)))
    
    
    print(sparse_name)
    print(np.mean(gflops_list))
    print(max(np.abs(gflops_list - init_gflops)))
    
    mean_gflops_list.append(np.mean(gflops_list))
    std_gflops_list.append(max(np.abs(gflops_list - init_gflops)))
    
print(sparse_name_list)
print(mean_gflops_list)
print(std_gflops_list)


'''

gen2
12.007653066666666
3.27768

ch7-8-b3
26.896418400000005
27.716386

3D_28984_Tetra
60.81071433333333
37.764226

'''

'''
['gen2', '176bit', 'bcsstk13', 'biplane-9', 
'ch7-8-b3', 'luxembourg_osm', 'cond-mat-2003', 'astro-ph', 
'3D_28984_Tetra', 'rajat25', 'c-69', 'rajat16', 
'net4-1', 'tsyl201', 'matrix-new_3', 'ct20stif']

[12.105580833333335, 11.402420866666663, 12.685011566666667, 12.8244018, 
21.847068166666663, 18.188226566666668, 27.9626386, 32.20923176666666, 
63.290969933333336, 34.6643146, 45.09779846666667, 34.86917496666667, 
70.82388256666667, 111.82776103333332, 95.45917230000002, 114.73698449999999]

[3.188319, 5.224155, 2.9092020000000005, 6.623354000000001, 
23.881994, 35.878761999999995, 14.316358000000005, 11.807281, 
39.29888999999999, 17.756554000000005, 13.180881000000007, 18.024985000000004, 
64.823845, 67.366406, 55.438053000000004, 56.516801]


['t2d_q9', 'fv2', 'fv3', 'big_dual', 
'wing', 'TSOPF_RS_b39_c7', 'airfoil_2d', 'jan99jac120sc', 
'Trec13', 'Zd_Jac2_db', 'bibd_17_8', 
'image_interp', 'stat96v2', 'belgium_osm', 'Chevron3', 'parabolic_fem']

[14.003952466666666, 13.505936233333333, 13.559925433333332, 12.817948266666667, 
20.1625665, 36.34767513333334, 37.919865699999995, 30.892226500000003, 
63.39158656666666, 57.67617016666666, 31.413409433333335, 28.61875613333333, 
92.9466017, 23.4599882, 80.09851566666664, 64.95922283333333]

[5.528461999999999, 5.830273999999999, 5.595079999999999, 9.971711000000003, 
20.551302, 14.371388999999997, 7.590357999999998, 20.52571, 
66.67237399999999, 27.860515999999997, 47.064499, 41.845529, 
50.81777700000001, 45.920216, 44.202059, 43.39931]
'''